#coding:utf-8
from __future__ import division
import os
from glob import glob
import random
from errors import TrainSamplesNotEnough
import logging
import traceback
import shutil


class Select(object):
    def __init__(self, class_name, train_num, pic_ext='jpg'):
        """
        train_num：要求的 训练数据 至少的 样本数
        """
        self.class_name = class_name
        self.train_num = train_num
        self.pic_ext = pic_ext

    def __get_all_pic_list(self, base_dir):
        class_base_dir = os.path.join(base_dir, self.class_name)
        all_pic_list = glob(os.path.join(class_base_dir, '*.' + self.pic_ext))
        return all_pic_list

    def get_train_pic_list(self, base_dir):
        all_pic_list = self.__get_all_pic_list(base_dir)
        random.shuffle(all_pic_list)
        if len(all_pic_list) < self.train_num:
            raise TrainSamplesNotEnough
        else:
            return all_pic_list[0:self.train_num]

    def get_split_pic_list(self, base_dir, train_split_factor):
        assert train_split_factor > 0 and train_split_factor < 1
        all_pic_list = self.__get_all_pic_list(base_dir)
        random.shuffle(all_pic_list)
        # print(len(all_pic_list))
        if len(all_pic_list) < int(self.train_num / train_split_factor):
            raise TrainSamplesNotEnough
        all_pic_list = all_pic_list[0: int(self.train_num / train_split_factor)]
        train_num = self.train_num
        train_pic_list = all_pic_list[0:train_num]
        validation_pic_list = all_pic_list[train_num:len(all_pic_list)]
        return train_pic_list, validation_pic_list


class OriginData(object):
    def __init__(self, base_dir, class_list, train_num):
        self.selects = []
        for class_name in class_list:
            self.selects.append(Select(class_name, train_num))
        self.base_dir = base_dir

    def create_split_dir(self, dir_name, train_split_factor):
        train_dir = os.path.join(dir_name, 'train')
        validation_dir = os.path.join(dir_name, 'validation')
        if os.path.exists(dir_name):
            choice = raw_input('%s already exists. delete? y/n: ' % dir_name)
            if choice == 'y':
                shutil.rmtree(dir_name)
            else:
                return
        if os.path.exists(dir_name) is False:
            os.makedirs(dir_name)
        if os.path.exists(train_dir) is False:
            os.makedirs(train_dir)
        if os.path.exists(validation_dir) is False:
            os.makedirs(validation_dir)
        for select in self.selects:
            try:
                train_list, validation_list = select.get_split_pic_list(self.base_dir, train_split_factor)
                class_train_dir = os.path.join(train_dir, select.class_name)
                class_validation_dir = os.path.join(validation_dir, select.class_name)
                os.makedirs(class_train_dir)
                os.makedirs(class_validation_dir)
                for filename in train_list:
                    shutil.copy(filename, class_train_dir)
                for filename in validation_list:
                    shutil.copy(filename, class_validation_dir)
                logging.info('Class %s finished.' % select.class_name)
            except TrainSamplesNotEnough:
                logging.error('picture samples of class "%s" are not enough!' % select.class_name)
            except Exception:
                logging.error(traceback.format_exc())
